#!/usr/bin/python3.8

import time
import numpy as np
from typing import List
from typing import Callable
from log import logger
from lib import xCoreSDK_python as robot
import time
Array6 = List[float]


class Arm:
    def init(self, **kwargs) -> bool:
        return False

    def close(self) -> None:
        pass

    def getState(self) -> int:
        return -1

    def getPose(self) -> Array6:
        return list(6 * (0,))

    def setPose(self, callback: Callable[[int], bool], *pose: Array6) -> bool:
        return False


class Rokae(Arm):
    r = None

    def init(self, **kwargs) -> bool:
        try:
            self.poses = []
            self.ec = {"ec": 0, "message": ""}
            ip = kwargs.get("ip", "192.168.2.160")

            self.r = robot.xMateRobot(ip)  # type: robot.xMateRobot
            self.ec["ec"] or self.r.connectToRobot(self.ec)
            self.ec["ec"] or self.r.recoverState(1, self.ec)

            self.ec["ec"] or self.r.setOperateMode(robot.OperateMode.automatic, self.ec)
            self.ec["ec"] or self.r.setPowerState(True, self.ec)
            self.ec["ec"] or self.r.adjustSpeedOnline(1, self.ec)
            self.ec["ec"] or self.r.setDefaultSpeed(1000, self.ec)
            self.ec["ec"] or self.r.setDefaultZone(0, self.ec)
            super().init(**kwargs)
        except Exception as e:
            logger.error(f"init arm failed with err: {e}")
            return False
        return not self.ec["ec"]

    def setAuto(self):
        self.r.recoverState(1, self.ec)
        self.r.setOperateMode(robot.OperateMode.automatic, self.ec)
        self.r.setPowerState(True, self.ec)
        self.r.adjustSpeedOnline(1, self.ec)
        self.r.setDefaultSpeed(1000, self.ec)
        self.r.setDefaultZone(0, self.ec)
        return self.r.powerState(self.ec) == robot.PowerState.on

    def close(self) -> None:
        # self.r.stop(self.ec)
        self.r.setPowerState(False, self.ec)
        self.r.disconnectFromRobot(self.ec)
        super().close()

    __del__ = close

    def getState(self) -> int:
        return int(self.r.operationState(self.ec))

    def drag(self, drag, tool="tool1"):
        self.r.setToolset(tool, "wobj0", self.ec)
        if drag:
            self.r.enableDrag(robot.DragParameterSpace.cartesianSpace.value, robot.DragParameterType.translationOnly.value, self.ec)
        else:
            self.r.disableDrag(self.ec)
        return bool(self.ec)

    def message(self) -> str:
        return self.ec.get("message", "")

    def getPose(self) -> Array6:
        pose = self.r.cartPosture(robot.CoordinateType.flangeInBase, self.ec)
        return pose.trans + pose.rpy

    def stop(self):
        self.r.stop(self.ec)
        self.r.moveReset(self.ec)

    def poseCheck(self, pose):
        cur = self.getPose()
        diff = np.sum(np.absolute(np.asarray(cur) - np.asarray(pose)))
        if diff > 0.01:
            logger.warning("{0} {1} {2}".format(diff, cur, pose))
            return False
        return True

    def pointsQuery(self, inv=False):
        res = self.r.queryEventInfo(robot.Event.moveExecution, self.ec)
        idx = res["wayPointIndex"] + int(res["reachTarget"])
        return self.poses[idx:] if inv else self.poses[:idx][::-1]

    def setPoseNonblock(self, sVel, vel, eVel, save=True, *poses: Array6) -> bool:
        self.r.moveReset(self.ec)
        size = len(poses) - 1
        cmd = [robot.MoveLCommand(p, sVel if idx == 0 else (eVel if idx == size else vel), 10) for idx, p in enumerate(poses)]
        self.r.executeCommand(cmd, self.ec)
        self.r.moveStart(self.ec)
        if save:
            self.poses = poses
        return True

    def setPose(self, callback: Callable[[int], bool], *poses: Array6) -> bool:
        return self.setPose2(callback, 200, 200, 200, *poses)

    def setPose2(self, callback: Callable[[int], bool], sVel, vel, eVel, *poses: Array6) -> bool:
        return self.setPose3(callback, sVel, vel, eVel, True, *poses)

    def setPose3(self, callback: Callable[[int], bool], sVel, vel, eVel, save, *poses: Array6) -> bool:
        self.setPoseNonblock(sVel, vel, eVel, save, *poses)
        while True:
            time.sleep(0.05)
            st = int(self.r.operationState(self.ec))
            if st in [-1, 0, 1]:
                self.r.stop(self.ec)
                break

            if not callback(st):
                self.r.stop(self.ec)
                return False

        return self.poseCheck(poses[-1])

    def setEndClose (self):
        ec = {}
        self.r.setDO(1,  0, True, ec)
        self.r.setDO(1,  1, True, ec)
        time.sleep(1.5)
        self.r.setDO(1,  0, False, ec)
        self.r.setDO(1,  1, False, ec)
        return ec

    def setEndOpen (self):
        ec = {}
        self.r.setDO(1,  0, False, ec)
        self.r.setDO(1,  1, True, ec)
        time.sleep(1.5)
        self.r.setDO(1,  0, False, ec)
        self.r.setDO(1,  1, False, ec)
        return ec


a = [0.2532188848287554, 0.09030179288176225, 0.3384137821651567, -3.1409692078286757, -0.00048430784375869286, 2.3215618536316596]
b =  [0.10680710828535048, -0.20382760394377142, 0.3384161440699347, -3.1409714661824824, -0.00048260724761573523, 0.816773405620627]

if __name__ == "__main__":
    arm = Rokae()
    arm.init()

    arm.setPose2(lambda st: True, 50, 50, 50, a,b)
    arm.setPose2(lambda st: True, 50, 50, 50, b,a)
    print(arm.getPose())
    time.sleep(1)
    arm.close()
    pass
